Skip to content

feat(zarr-indexing): factor chunk plans into a columnar GridPartition - #4310

Merged
d-v-b merged 20 commits into
zarr-developers:mainfrom
d-v-b:zarr-indexing/grid-partition
Sep 6, 2026
Merged

feat(zarr-indexing): factor chunk plans into a columnar GridPartition#4310
d-v-b merged 20 commits into
zarr-developers:mainfrom
d-v-b:zarr-indexing/grid-partition

Conversation

@d-v-b

@d-v-b d-v-b commented Sep 3, 2026

Copy link
Copy Markdown
Contributor

Summary

This is an update to the data structures used in zarr-indexing for relating a selection on a chunked array to a per-chunk plan. We currently generate objects per chunk eagerly, which is inefficient due to object creation overhead. FWIW I investigated pushing this down to rust but crossing the py03 boundary for a ton of tiny objects ends up pretty expensive.

a better solution is to model the chunk plan as the result of a product of per-axis plans, which is what this PR adds. the end result restores a lot of performance that had been lost with the eager object creation.

based on a claude-authored PR here: d-v-b#316

Author attestation

  • I am a human, these are my changes, and I have reviewed and understood every change and can explain why each is correct.

TODO

  • Add unit tests and/or doctests in docstrings
  • Add docstrings and API docs for any new/modified user-facing classes and functions
  • New/modified features documented in docs/user-guide/*.md
  • Changes documented as a new file in changes/
  • GitHub Actions have all passed
  • Test coverage is 100% (Codecov passes)

Restricting a transform to a chunk box distributes over output dimensions
whenever each output map reads its own input axis, which is every basic
and orthogonal selection. Chunk resolution therefore no longer intersects
the whole transform with every candidate chunk; it resolves each axis once
against its grid into a table (StridedSet / IndexedSet), sorts correlated
(vindex) index arrays into chunks once into a JointSet, and derives each
ChunkProjection as one row of each table. ChunkPlan.partition() and
partition_transform() expose the factored form, so a consumer can read
the tables directly instead of materializing an object graph per chunk.

The projections a plan yields are unchanged; the general whole-transform
walk remains for hand-built diagonals, which have no factored form.

Along the way: _intersect_general reuses a precomputed _CorrelatedBlock and
accepts survivor positions; checked_affine has identity and dtype-bounded
fast paths; ArrayMap._with_affine shares frozen index arrays on translate;
IndexDomain._unchecked / IndexTransform._unchecked skip validation for
objects derived from an already-valid transform.

Assisted-by: ClaudeCode:claude-fable-5-1
Assisted-by: ClaudeCode:claude-fable-5-1
The package and its tests import nothing from zarr; the old comments claimed
the chunk-resolution tests needed zarr's ChunkGrid, which stopped being true
once the package grew its own grids. The real reason is the shared pinned
test toolchain.

Assisted-by: ClaudeCode:claude-fable-5-1
…hunk narrative

The module docstring described intersecting the whole transform with every
candidate chunk as "the algorithm"; that walk is now the fallback for
hand-built diagonals only. It now explains the factored form and its three
tables, and why they cost the sum of the touched chunks per axis.

The visual guide gains a final integrator section, "A plan is a product of
per-axis tables", with an executable snippet that reads the StridedSet,
IndexedSet and JointSet tables off real plans and checks the plan's
projections against the partition's rows. Integration boundaries gains
"Reading the tables directly", a consumer that assembles a strided box from
the tables with no projection materialized. The API index, landing page and
design notes (TensorStore lineage, the performance caveat, and the box/query
split) point at the new section.

Assisted-by: ClaudeCode:claude-fable-5-1
…w fixes

Adversarial review (roborev, a correctness reviewer, a complexity reviewer,
and ~24k differential examples against main) of the grid partition.

Cuts. The whole-transform walk that remained for hand-built diagonals is
gone: it was unreachable for every index-array shape, its key builder was
duplicated verbatim in _chunk_keys, and for the one shape it served it
produced wrong projections (a three-point diagonal yielded four projections
covering six cells, on main too). A DimensionMap diagonal is now rejected
with ValueError. With it go the sorted-1-D fast path, the three cell-transform
helpers, the block/positions parameters of _intersect_general, the
correlated-residual check that admitted a diagonal and then crashed,
GridPartition.__getitem__, partition_transform as public API, the
object-dtype column fallback (StridedSet.origin is now a position along the
request axis, so every column is intp), checked_affine's dtype-bound
shortcut (measured at noise; the identity shortcut stays and now accepts
bool via np.can_cast, as main did), and StridedSet.chunk_map/cell_map.

Fixes. GridPartition.n_rows is an exact integer and len raises OverflowError
instead of wrapping to zero; table columns are read-only, so a memoized
partition cannot drift under a consumer; the documented table consumer now
handles reversed axes, inserted axes and transposed transforms, and the
snippet checks all three.

Docs. Corrected the diagonal statement everywhere it appeared, the memoized
"fresh walk" wording, the "vectorized per axis" claim, and the TensorStore
correspondence (its strided sets are per input dimension; it keeps one index
array set per connected component). The guide no longer restates the class
docstrings.

Assisted-by: ClaudeCode:claude-fable-5-1
A zero-stride DimensionMap over a domain wider than np.intp is valid and
touches one storage cell; coercing every StridedSet column to intp made it
raise OverflowError where main returned one projection. `extent` and
`origin` are the two columns measured along the request axis, whose bounds
are arbitrary Python ints, so they now fall back to exact-int (object)
columns when a value does not fit. Chunk-local columns stay intp.

Also corrects the design note that said both affine-diagonal cases raise
NotImplementedError: two slice maps sharing an axis now raise ValueError.

Assisted-by: ClaudeCode:claude-fable-5-1
@github-actions github-actions Bot added the needs release notes Automatically applied to PRs which haven't added release notes label Sep 3, 2026
@read-the-docs-community

read-the-docs-community Bot commented Sep 3, 2026

Copy link
Copy Markdown

@codecov

codecov Bot commented Sep 3, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.23%. Comparing base (a50d791) to head (da30a49).

Additional details and impacted files
@@           Coverage Diff           @@
##             main    #4310   +/-   ##
=======================================
  Coverage   94.23%   94.23%           
=======================================
  Files          92       92           
  Lines       12880    12880           
=======================================
  Hits        12137    12137           
  Misses        743      743           
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

…d the minimal grid protocol

Every varying grid in the partition cases summed exactly to its extent, so
the boundary where a chunk's data extent is shorter than its declared size,
the rectilinear-specific case, was unpinned; so was a grid without
data_size. Both now run through the evaluation oracle for strided,
orthogonal and correlated selections.

Assisted-by: ClaudeCode:claude-fable-5-1
@d-v-b
d-v-b marked this pull request as ready for review September 3, 2026 18:22
d-v-b and others added 12 commits September 3, 2026 20:22
Retain public selection-flow documentation and fix singleton data-extent coverage. Move the execution prototype to a follow-up review.

Assisted-by: Codex:GPT-6
… columns, one planning mechanism

Adversarial review of the branch (roborev, a correctness reviewer, a
complexity reviewer, ~40k differential examples against main and NumPy).

Fixes. Correlated planning probed storage bounds only through the grid's
vectorized lookup, which zarr's grids do not validate, so an out-of-range
coordinate silently planned a chunk that does not exist; the joint table now
probes each component's extreme coordinates with the scalar lookup, as the
orthogonal table already did. Table columns were read-only by flag only,
which setflags(write=True) undoes; they are now re-homed over immutable
bytes, like ArrayMap's index array, so a memoized partition cannot drift.

Cuts, each with what it cost stated in the review: the affine-diagonal
iteration path (a second mechanism that reintroduced transform.intersect
per chunk for a shape no selection produces; diagonals now raise ValueError
from iteration as from partition()), the single-component iterator (14-23%
on a walk that stays 2-3x behind zarr's coordinate indexer either way),
chunk_coord_batches with its beyond-intp mixed-radix path and n_rows (no
consumer; a partition with row_shape writes the batching in three lines),
and the stored block_coordinates column (byte-identical to positions on
every 1-D block; now a memoized property). Tests that pinned memoization
identity and column flags collapse into one that asserts immutability.

Docs: the diagonal statement is now precise everywhere it appears,
StridedSet.full no longer claims "in order", the guide no longer says sets
holds every source axis, and the changelog describes the feature as shipped
rather than the branch's history.

Assisted-by: ClaudeCode:claude-fable-5-1
…ng to intp

A uint64 value beyond the intp range wrapped to a negative index on the
cast and then passed as a wrapped position, selecting the last element
instead of raising. Unsigned values are never negative, so they are
checked as they are and narrowed only once in range.

Assisted-by: ClaudeCode:claude-fable-5-1
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
@d-v-b

d-v-b commented Sep 6, 2026

Copy link
Copy Markdown
Contributor Author

this got a lot of review / improvement from claude and codex. we are pretty close to zarr's current indexing performance on the eager path. merging.

@d-v-b
d-v-b merged commit ef6cfe2 into zarr-developers:main Sep 6, 2026
41 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

needs release notes Automatically applied to PRs which haven't added release notes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant